{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import argparse\n",
    "import cv2\n",
    "import numpy as np\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description=\"Homography from Source to Destination Image\")\n",
    "parser.add_argument('-id', '--CAMERA_ID', default=1, type=int, help='Camera ID')\n",
    "parser.add_argument('-path', '--INPUT_PATH', default='./data/', type=str, help='Input Source/Destination Image Path')\n",
    "parser.add_argument('-bw','--BORAD_WIDTH', default=7, type=int, help='Chess Board Width (corners number)')\n",
    "parser.add_argument('-bh','--BORAD_HEIGHT', default=6, type=int, help='Chess Board Height (corners number)')\n",
    "parser.add_argument('-src', '--SOURCE_IMAGE', default='img_src', type=str, help='Source Image File Name Prefix (eg.:img_src)')\n",
    "parser.add_argument('-dst', '--DEST_IMAGE', default='img_dst', type=str, help='Destionation Image File Name Prefix (eg.:img_dst)')\n",
    "parser.add_argument('-size','--SCALED_SIZE', default=10, type=int, help='Scaled Chess Board Square Size (image pixel)')\n",
    "parser.add_argument('-subpix_s','--SUBPIX_REGION_SRC', default=3, type=int, help='Corners Subpix Region of img_src')\n",
    "parser.add_argument('-subpix_d','--SUBPIX_REGION_DST', default=3, type=int, help='Corners Subpix Region of img_dst')\n",
    "parser.add_argument('-center','--CENTER_FLAG', default=False, type=bool, help='Center Image Manually (Ture/False)')\n",
    "parser.add_argument('-scale','--SCALE_FLAG', default=False, type=bool, help='Scale Image to Fix Board Size (Ture/False)')\n",
    "parser.add_argument('-store','--STORE_FLAG', default=False, type=bool, help='Store Centerd/Scaled Images (Ture/False)')\n",
    "parser.add_argument('-store_path', '--STORE_PATH', default='./data/', type=str, help='Path to Store Centerd/Scaled Images')\n",
    "args = parser.parse_args([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# args.CAMERA_ID = 1                         # 相机编号\n",
    "# args.INPUT_PATH = './data/'                # 图片、视频输入路径\n",
    "# args.BORAD_WIDTH = 7                       # 棋盘宽度 【内角点数】\n",
    "# args.BORAD_HEIGHT = 6                      # 棋盘高度 【内角点数】\n",
    "# args.SOURCE_IMAGE = 'img_src'              # 源图像文件名前缀\n",
    "# args.DEST_IMAGE = 'img_dst'                # 目标图像文件名前缀\n",
    "# args.SQUARE_SIZE = 10                      # 缩放后棋盘每小格边长占据图中的像素尺寸\n",
    "# args.SUBPIX_REGION_SRC = 3                 # 源图像亚像素优化的搜索像素范围(根据图像分辨率调整)\n",
    "# args.SUBPIX_REGION_DST = 3                 # 目标图像亚像素优化的搜索像素范围(根据图像分辨率调整)\n",
    "# args.CENTER_FLAG = True                    # 是否对目标图像进行居中\n",
    "# args.SCALE_FLAG = True                     # 是否对目标图像进行缩放\n",
    "# args.STORE_FLAG = True                     # 是否储存目标图像（居中缩放后）\n",
    "# args.STORE_PATH = './data/'                # 目标图像（居中缩放后）保存路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class CenterImage:        # 鼠标点击位置平移至图像中心,使得车辆中心位于图像中心\n",
    "    def __init__(self):\n",
    "        self.x = 0\n",
    "        self.y = 0\n",
    "        self.param = {'tl': None, 'br': None, 'current_pos': None,'complete': False}  # topleft bottomright\n",
    "        self.display = \"CLICK image center and press Y/N to validate, ESC to stay original\"\n",
    "    \n",
    "    # 鼠标框选返回中心点模式 或直接点击中心点\n",
    "    def mouse(self, event, x, y, flags, param):\n",
    "        if event == cv2.EVENT_LBUTTONDOWN:      # 鼠标左键按下\n",
    "            img = self.raw_frame.copy()\n",
    "            param['current_pos'] = (x, y)\n",
    "            if param['tl'] is None:\n",
    "                param['tl'] = param['current_pos'] \n",
    "        if event == cv2.EVENT_MOUSEMOVE and param['tl'] is not None and not param['complete']: # 鼠标移动\n",
    "            img = self.raw_frame.copy()\n",
    "            param['current_pos'] = (x, y)\n",
    "            cv2.rectangle(img, param['tl'], param['current_pos'], (0, 0, 255))\n",
    "            cv2.imshow(self.display, img)\n",
    "        if event == cv2.EVENT_LBUTTONUP and param['tl'] is not None:   # 鼠标左键抬起 完成绘制\n",
    "            img = self.raw_frame.copy()\n",
    "            param['br'] = (x, y)\n",
    "            param['complete'] = True\n",
    "            cv2.rectangle(img, param['tl'], param['br'], (0, 0, 255))\n",
    "            cv2.imshow(self.display, img)\n",
    "            self.x = (param['tl'][0] + param['br'][0] ) // 2\n",
    "            self.y = (param['tl'][1] + param['br'][1] ) // 2\n",
    "            text = \" %d,%d? (y/n)\" % (self.x, self.y)\n",
    "            cv2.circle(img, (self.x, self.y), 1, (0, 0, 255), thickness = 2)\n",
    "            cv2.putText(img, text, (self.x, self.y), cv2.FONT_HERSHEY_PLAIN,1.0, (0, 0, 0), thickness = 1)\n",
    "            cv2.imshow(self.display, img)\n",
    "        self.param = param\n",
    "        \n",
    "    def translate(self, img):\n",
    "        shift_x = img.shape[1] // 2 - self.x\n",
    "        shift_y = img.shape[0] // 2 - self.y\n",
    "        M = np.float32([[1,0,shift_x],[0,1,shift_y]])\n",
    "        img_dst = cv2.warpAffine(img,M,img.shape[1::-1])\n",
    "        return img_dst\n",
    "        \n",
    "    def __call__(self, raw_frame):   \n",
    "        self.raw_frame = raw_frame\n",
    "        cv2.namedWindow(self.display, flags = cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)\n",
    "        cv2.setMouseCallback(self.display, self.mouse, self.param)\n",
    "        while True:\n",
    "            cv2.imshow(self.display, self.raw_frame)\n",
    "            key = cv2.waitKey(0)\n",
    "            if key in (ord(\"y\"), ord(\"Y\")):\n",
    "                break\n",
    "            elif key in (ord(\"n\"), ord(\"N\")):\n",
    "                self.x = 0\n",
    "                self.y = 0\n",
    "                self.param['tl'] = None\n",
    "                self.param['br'] = None\n",
    "                self.param['current_pos'] = None\n",
    "                self.param['complete'] = None\n",
    "            elif key == 27: \n",
    "                self.x = 0\n",
    "                self.y = 0\n",
    "                break\n",
    "        cv2.destroyAllWindows()\n",
    "        if not (self.x == 0 and self.y == 0):\n",
    "            return self.translate(self.raw_frame)\n",
    "        else:\n",
    "            return self.raw_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class ScaleImage:        # 缩放图像使得标定板在图中的像素尺寸保持一致\n",
    "    def __init__(self, corners):        \n",
    "        self.calc_dist(corners)\n",
    "        print('scale image from {} to {}'.format(self.dist_square,args.SCALED_SIZE))\n",
    "        self.scale_factor = args.SCALED_SIZE / self.dist_square      # 缩放系数\n",
    "\n",
    "    # 计算照片上棋盘每格的距离\n",
    "    def calc_dist(self, corners):\n",
    "        dist_total = 0\n",
    "        for i in range(args.BORAD_HEIGHT):\n",
    "            dist = cv2.norm(corners[i * args.BORAD_WIDTH,:], corners[(i+1) * args.BORAD_WIDTH-1,:], cv2.NORM_L2)\n",
    "            dist_total += dist / (args.BORAD_WIDTH - 1)\n",
    "        self.dist_square = dist_total / args.BORAD_HEIGHT\n",
    "        \n",
    "    # 图像填充黑边\n",
    "    def padding(self, img, width, height):\n",
    "        H = img.shape[0]\n",
    "        W = img.shape[1]\n",
    "        top = (height - H) // 2 \n",
    "        bottom = (height - H) // 2 \n",
    "        if top + bottom + H < height:\n",
    "            bottom += 1\n",
    "        left = (width - W) // 2 \n",
    "        right = (width - W) // 2 \n",
    "        if left + right + W < width:\n",
    "            right += 1\n",
    "        return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value = (0,0,0))  \n",
    "    \n",
    "    # 图像裁剪 \n",
    "    def crop(self, img, width, height):\n",
    "        H = img.shape[0]\n",
    "        W = img.shape[1]\n",
    "        top = (H - height) // 2\n",
    "        bottom = (H - height) // 2 + height\n",
    "        left = (W - width) // 2\n",
    "        right = (W - width) // 2 + width\n",
    "        return img[top:bottom, left:right]          \n",
    "    \n",
    "    # 进行图像缩放及之后的填充或裁剪操作\n",
    "    def __call__(self, raw_frame):\n",
    "        width = raw_frame.shape[1]\n",
    "        height = raw_frame.shape[0]\n",
    "        raw_frame = cv2.resize(raw_frame, (0,0), fx=self.scale_factor, fy=self.scale_factor)  # 图像缩放\n",
    "        if self.scale_factor < 1:\n",
    "            raw_frame = self.padding(raw_frame, width, height)\n",
    "        else:                     \n",
    "            raw_frame = self.crop(raw_frame, width, height)\n",
    "        return raw_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExCalibrator():           # 外参标定器，从两张去畸变图像得到单应性矩阵\n",
    "    def __init__(self):\n",
    "        self.src_corners_total = np.empty([0,1,2])\n",
    "        self.dst_corners_total = np.empty([0,1,2])\n",
    "    \n",
    "    # 获取args参数，供外部调用修改参数\n",
    "    @staticmethod\n",
    "    def get_args():\n",
    "        return args\n",
    "    \n",
    "    def imgPreprocess(self, img, center, scale):\n",
    "        # 俯视图图像居中\n",
    "        if center:\n",
    "            centerImg = CenterImage()\n",
    "            img = centerImg(img)\n",
    "        # 图像尺度缩放\n",
    "        if scale:\n",
    "            ok, corners = self.get_corners(img, subpix = args.SUBPIX_REGION_DST)\n",
    "            if not ok:\n",
    "                raise Exception(\"failed to find corners in destination image\")\n",
    "            scaleImg = ScaleImage(corners)\n",
    "            img = scaleImg(img)\n",
    "        cv2.imshow(\"Preprocessed Image\", img)\n",
    "        cv2.waitKey(0)\n",
    "        return img\n",
    "        \n",
    "    def get_corners(self, img, subpix, draw=False):\n",
    "        ok, corners = cv2.findChessboardCorners(img, (args.BORAD_WIDTH, args.BORAD_HEIGHT),\n",
    "                      flags = cv2.CALIB_CB_ADAPTIVE_THRESH|cv2.CALIB_CB_NORMALIZE_IMAGE)\n",
    "        if ok: \n",
    "            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "            # 角点坐标亚像素优化\n",
    "            corners = cv2.cornerSubPix(gray, corners, (subpix, subpix), (-1, -1),\n",
    "                                       (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.01))\n",
    "        if draw:\n",
    "            cv2.drawChessboardCorners(img, (args.BORAD_WIDTH, args.BORAD_HEIGHT), corners, ok)   # 在图上画出角点\n",
    "        return ok, corners\n",
    "    \n",
    "    def warp(self):\n",
    "        src_warp = cv2.warpPerspective(self.src_img, self.homography, \n",
    "                                       (self.dst_img.shape[1], self.dst_img.shape[0])) \n",
    "        return src_warp\n",
    "        \n",
    "    def __call__(self, src_img, dst_img):\n",
    "        ok, dst_corners = self.get_corners(dst_img, subpix = args.SUBPIX_REGION_DST, draw=True)\n",
    "        if not ok:\n",
    "            raise Exception(\"failed to find corners in destination image\")\n",
    "        ok, src_corners = self.get_corners(src_img, subpix = args.SUBPIX_REGION_SRC, draw=True)\n",
    "        if not ok:\n",
    "            raise Exception(\"failed to find corners in source image\")\n",
    "        self.dst_corners_total = np.append(self.dst_corners_total, dst_corners, axis = 0)\n",
    "        self.src_corners_total = np.append(self.src_corners_total, src_corners, axis = 0)\n",
    "        self.homography, _ = cv2.findHomography(self.src_corners_total, self.dst_corners_total,method = cv2.RANSAC)\n",
    "        self.src_img = src_img\n",
    "        self.dst_img = dst_img\n",
    "        return self.homography"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_images(PATH, NAME):\n",
    "        filePath = [os.path.join(PATH, x) for x in os.listdir(PATH) \n",
    "                    if any(x.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])\n",
    "                   ]                                                           # 得到给定路径下所有图片文件\n",
    "        filenames = [filename for filename in filePath if NAME in filename]    # 再筛选出包含给定名字的图片\n",
    "        if len(filenames) == 0:\n",
    "            raise Exception(\"from {} read images failed\".format(PATH))\n",
    "        return filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def main():\n",
    "    # 去畸变后的图像输入 \n",
    "    srcfiles = get_images(args.INPUT_PATH, args.SOURCE_IMAGE)\n",
    "    dstfiles = get_images(args.INPUT_PATH, args.DEST_IMAGE)  \n",
    "    if len(srcfiles) != len(dstfiles):\n",
    "        raise Exception(\"numbers of source and destination images should be equal\")\n",
    "    \n",
    "    exCalib = ExCalibrator()                 # 初始化外参标定器\n",
    "\n",
    "    for i in range(len(srcfiles)):    \n",
    "        src_raw = cv2.imread(srcfiles[i])\n",
    "        dst_raw = cv2.imread(dstfiles[i])\n",
    "        print(srcfiles[i])\n",
    "        print(dstfiles[i])\n",
    "        \n",
    "        # 需要时使用imgPreprocess方法对destination图像进行居中和缩放\n",
    "        if args.CENTER_FLAG or args.SCALE_FLAG:\n",
    "            dst_raw = exCalib.imgPreprocess(dst_raw, args.CENTER_FLAG, args.SCALE_FLAG)\n",
    "        if args.STORE_FLAG:\n",
    "            cv2.imwrite(args.STORE_PATH + 'img_dst{}.jpg'.format(i), dst_raw)  \n",
    "        \n",
    "        # 输入对应两张去畸变图像得到单应性矩阵\n",
    "        homography = exCalib(src_raw, dst_raw)\n",
    "        np.save('camera_{}_H.npy'.format(args.CAMERA_ID), homography)\n",
    "        \n",
    "        # 得到source图像的变换结果图\n",
    "        src_warp = exCalib.warp()\n",
    "        \n",
    "        cv2.namedWindow(\"Source View\", flags = cv2.WINDOW_NORMAL|cv2.WINDOW_KEEPRATIO)\n",
    "        cv2.imshow(\"Source View\", src_raw)\n",
    "        cv2.namedWindow(\"Destination View\", flags = cv2.WINDOW_NORMAL|cv2.WINDOW_KEEPRATIO)\n",
    "        cv2.imshow(\"Destination View\", dst_raw)\n",
    "        cv2.namedWindow(\"Warped Source View\", flags = cv2.WINDOW_NORMAL|cv2.WINDOW_KEEPRATIO)\n",
    "        cv2.imshow(\"Warped Source View\", src_warp)\n",
    "        \n",
    "        while True:\n",
    "            key = cv2.waitKey(0)\n",
    "            if key == 27: break         # ESC退出\n",
    "        cv2.destroyAllWindows()\n",
    "    \n",
    "           \n",
    "if __name__ == '__main__':\n",
    "    main()    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
